from tqdm import tqdm
import math
import numpy as np
import ROOT as rt

rt.gROOT.ProcessLine("struct TreeVar_t {\
    Float_t         Theta;\
    Float_t         Phi;\
    Float_t         pTheta;\
    Float_t         pPhi;\
    Float_t         Weight;\
};" )

nMCBins = int(40)
nDalitzBins = int(200)
nDalitzMeerkatBins = int(600)
nDivBins = int(20)

kernelWidth = 0.05
sampleBins = 200

thetaRange = [2.5, np.pi]
pthetaRange = [0., np.pi]
phiRange = [-np.pi, np.pi]

def GenRoot():
    TreeVar = rt.TreeVar_t()
    InTree = rt.TChain("MCData")
    InTree.Add("GammaInDME.root")
    InTree.Add("GammaInDMEEnlarge.root")
    NumOfEntries = InTree.GetEntries()

    OutFile = rt.TFile("GammaThetaAndPhi.root", "recreate")
    OutTree = rt.TTree("MCData", "")
    OutTree.Branch("Theta", rt.addressof(TreeVar, "Theta"), "Theta/F")
    OutTree.Branch("Phi", rt.addressof(TreeVar, "Phi"), "Phi/F")
    OutTree.Branch("pTheta", rt.addressof(TreeVar, "pTheta"), "pTheta/F")
    OutTree.Branch("pPhi", rt.addressof(TreeVar, "pPhi"), "pPhi/F")
    OutTree.Branch("Weight", rt.addressof(TreeVar, "Weight"), "Weight/F")

    for entry in tqdm(InTree, total=NumOfEntries):
        vec = rt.TVector3(entry.PassX, entry.PassY, entry.PassZ)
        TreeVar.Theta = vec.Theta()
        TreeVar.Phi = vec.Phi()
        pvec = rt.TVector3(entry.PassPX, entry.PassPY, entry.PassPZ)
        TreeVar.pTheta = pvec.Theta()
        TreeVar.pPhi = pvec.Phi()
        TreeVar.Weight = 1.
        OutTree.Fill()

    OutFile.cd()
    OutTree.AutoSave()
    OutFile.Close()

def GenMap():

    import os, sys
    os.environ["ROOT_INCLUDE_PATH"] = "$MEERKATROOT/inc/"
    rt.gSystem.Load("$MEERKATROOT/lib/libMeerkat.so")

    InFile = rt.TFile(f"GammaThetaAndPhi.root")
    InTree = InFile.Get("MCData")

    thetaPHSP  = rt.OneDimPhaseSpace("thetaPHSP" , thetaRange[0], thetaRange[1])
    phiPHSP = rt.OneDimPhaseSpace("phiPHSP", phiRange[0], phiRange[1])
    twoDimPHSP = rt.CombinedPhaseSpace("twoDimPHSP", thetaPHSP, phiPHSP)

    GammaPosition = rt.TH2F("GammaPosition","Dalitz plot",nDalitzMeerkatBins, thetaRange[0], thetaRange[1], nDalitzMeerkatBins, phiRange[0], phiRange[1])
    GammaPosition.GetXaxis().SetTitle("Theta")
    GammaPosition.GetYaxis().SetTitle("Phi")

    KDE_gen = rt.BinnedKernelDensity(
            "KernelPDF",
            twoDimPHSP,
            InTree,
            "Theta", "Phi",
            "Weight",
            sampleBins, sampleBins,
            kernelWidth, kernelWidth,
            0,
            10000000
        )

    OutFile = rt.TFile("GammaAngular.root", "recreate")
    OutFile.cd()
    KDE_gen.project(GammaPosition)

    pthetaPHSP  = rt.OneDimPhaseSpace("pthetaPHSP" , pthetaRange[0], pthetaRange[1])
    pphiPHSP = rt.OneDimPhaseSpace("pphiPHSP", phiRange[0], phiRange[1])
    ptwoDimPHSP = rt.CombinedPhaseSpace("ptwoDimPHSP", pthetaPHSP, pphiPHSP)

    GammaMomDir = rt.TH2F("GammaMomDir","Dalitz plot",nDalitzMeerkatBins, pthetaRange[0], pthetaRange[1], nDalitzMeerkatBins, phiRange[0], phiRange[1])
    GammaMomDir.GetXaxis().SetTitle("Theta")
    GammaMomDir.GetYaxis().SetTitle("Phi")

    pKDE_gen = rt.BinnedKernelDensity(
            "pKernelPDF",
            ptwoDimPHSP,
            InTree,
            "pTheta", "pPhi",
            "Weight",
            sampleBins, sampleBins,
            kernelWidth, kernelWidth,
            0,
            10000000
        )

    pKDE_gen.project(GammaMomDir)

    GammaPosition.Write("GammaPosition")
    GammaMomDir.Write("GammaMomDir")

    InFile.Close()
    OutFile.Close()

if __name__ == "__main__":
        #GenRoot()
        GenMap()
